from datasets import SiameseNetworkDataset
from siamese import SiameseNetwork
# 展示图片
def imshow(img, text=None, should_save=False):
    npimg = img.numpy()
    plt.axis("off")
    if text:
        plt.text(75,8,text, 
                style='italic', 
                fontweight='bold', 
                bbox={'facecolor': 'white', 'alpha': 0.8, 'pad': 10})
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()
# 参数设置, 加载数据集
testing_dir = r"/data/Siamese_for_Face/faces/testing"
input_shape = [224, 224]
transform = transforms.Compose([transforms.Resize((224, 224)),
                                      transforms.ToTensor()])
folder_dataset_test = dset.ImageFolder(testing_dir)
siamese_dataset = SiameseNetworkDataset( imageFolderDataset=folder_dataset_test,
                                                 transform=transform,
                                                 should_invert=False) 
test_dataloader = DataLoader(siamese_dataset,num_workers=0,batch_size=1,shuffle=True) 
dataiter = iter(test_dataloader)
x0, _, _ = next(dataiter)
# 加载网络和训练过的权重
net = SiameseNetwork(input_shape)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
net.load_state_dict(torch.load(r'weights/vgg_siamese.pt'))
# 测试过程
if __name__ == '__main__':
    for i in range(10):
        _, x1, label2 = next(dataiter)
        x0, x1, label2 = x0.to(device),x1.to(device),label2.to(device) 
        concatenated = torch.cat((x0, x1), 0)
        output = net(Variable(x0), Variable(x1))[0]
        output = torch.nn.Sigmoid()(output) 
        from torchvision.utils import make_grid
        imshow(make_grid(concatenated).cpu(),'similarity: {:.2f}'.format(output.item()))
